# -*- coding: utf-8 -*-
# @Time    : 2024/7/3 14:54
# @Author  : yujiahao
# @File    : 04_fatch_mysql.py
# @description: 该脚本演示如何使用数据库连接池从MySQL数据库批量读取和写入数据，并使用装饰器测量函数执行时间。

import pymysql
from dbutils.pooled_db import PooledDB
import logging
import time
from faker import Faker
import sys

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s', force=True)


def measure_time(func):
    """
    装饰器函数，用于测量函数执行时间。

    Args:
        func (function): 被装饰的函数。

    Returns:
        function: 包装后的函数。
    """

    def wrapper(*args, **kwargs):
        # 记录函数开始执行的时间
        start_time = time.time()
        logging.info(f"开始执行函数 {func.__name__}")

        # 调用被装饰的函数，并获取其返回值
        result = func(*args, **kwargs)

        # 检查返回值是否是可迭代对象（但不是字符串或字节序列）
        if hasattr(result, '__iter__') and not isinstance(result, (str, bytes)):
            def generator_wrapper(*args, **kwargs):
                # 逐项迭代返回值
                for item in result:
                    yield item
                # 记录函数结束执行的时间，并计算执行时间
                end_time = time.time()
                logging.info(f"函数 {func.__name__} 执行时间: {end_time - start_time:.2f} 秒")

            # 返回包装后的生成器
            return generator_wrapper()
        else:
            # 如果返回值不是可迭代对象，直接记录执行时间
            end_time = time.time()
            logging.info(f"函数 {func.__name__} 执行时间: {end_time - start_time:.2f} 秒")
            return result

    return wrapper


# 配置数据库连接池
def init_db_pool():
    """
    初始化数据库连接池。

    Returns:
        PooledDB: 数据库连接池对象。
    """
    pool = PooledDB(
        creator=pymysql,  # 创建数据库连接的模块
        maxconnections=20,  # 连接池允许的最大连接数
        mincached=5,  # 初始化时创建的连接数
        maxcached=10,  # 连接池中最多闲置的连接数
        blocking=True,  # 连接池中没有可用连接后是否阻塞等待
        maxusage=None,  # 单个连接的最大重复使用次数
        setsession=['SET AUTOCOMMIT = 1'],  # 会话参数
        ping=1,  # 检查连接是否可用
        host='127.0.0.1',  # 数据库主机地址
        port=33061,  # 数据库端口
        user='***',
        passwd='***',
        db='bigdata_test',  # 数据库名称
        charset='utf8'  # 数据库编码
    )
    return pool


# 初始化数据库连接池
db_pool = init_db_pool()


@measure_time
def batch_read(read_sql, batch_size=10000):
    """
    从数据库中按批次读取数据。

    Args:
        read_sql (str): 要执行的 SQL 查询语句。
        batch_size (int, optional): 每次获取的行数。默认值为10000。

    Yields:
        tuple: 读取的数据行。
    """
    logging.info("开始批量读取")
    conn = db_pool.connection()  # 获取数据库连接

    try:
        with conn.cursor() as cursor:
            cursor.execute(read_sql)
            logging.info(f"执行的 SQL 查询: {read_sql}")
            while True:
                rows = cursor.fetchmany(batch_size)
                if not rows:
                    break
                logging.info(f"读取到的行数: {len(rows)}")
                for row in rows:
                    yield row  # 使用生成器逐行返回数据
    except pymysql.MySQLError as e:
        logging.error(f"数据库错误: {e}")
    finally:
        conn.close()  # 关闭数据库连接

    logging.info("完成批量读取")


@measure_time
def batch_write(write_sql, data, batch_size=10000):
    """
    将数据按批次写入数据库。

    Args:
        write_sql (str): 要执行的 SQL 插入语句。
        data (list): 要写入的数据行列表。
        batch_size (int, optional): 每次写入的行数。默认值为10000。
    """
    logging.info("开始批量写入")
    conn = db_pool.connection()  # 获取数据库连接

    try:
        with conn.cursor() as cursor:
            for i in range(0, len(data), batch_size):
                batch = data[i:i + batch_size]
                values = ', '.join(str(tuple(row)) for row in batch)
                sql = write_sql.format(values=values)
                cursor.execute(sql)
            conn.commit()
    finally:
        conn.close()  # 关闭数据库连接

    logging.info("完成批量写入")


def generate_sample_data(num_records=10000):
    """
    生成示例数据。

    Args:
        num_records (int, optional): 要生成的记录数。默认值为10000。

    Returns:
        list: 生成的数据行列表。
    """
    fake = Faker()
    data = []

    for _ in range(num_records):
        name = fake.name()
        age = fake.random_int(min=18, max=65)
        department = fake.job()
        data.append((name, age, department))

    return data


def main():
    # 示例数据
    write_sql = "INSERT INTO bigdata_test.yujiahao (name, age, department) VALUES {values}"

    # 生成示例数据并插入目标表
    sample_data = generate_sample_data(10000)
    batch_write(write_sql, sample_data)

    # 批量读取示例
    read_sql = "SELECT * FROM bigdata_test.user_app_test limit 10000"
    for row in batch_read(read_sql):
        print(row)  # 逐行打印读取的数据
        sys.stdout.flush()  # 确保立即输出


if __name__ == "__main__":
    main()
